import paramiko
from paramiko import SSHException, AuthenticationException, BadHostKeyException
import logging
from concurrent.futures import ThreadPoolExecutor, as_completed

# 设置日志记录
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class SSHClient:
    '''
    初始化SSH连接参数

    参数：
        hostname：远程主机的名称或IP地址。
        username：用于登录远程主机的用户名。
        password：用于登录远程主机的密码。
        port：SSH连接的端口号，默认值为22。
        timeout：SSH连接的超时时间（秒），默认值为10。

    返回：
        无

    注意：
        该方法用于初始化SSH客户端对象的属性。
    '''
    def __init__(self, hostname, username, password, port=22, timeout=10):
        self.hostname = hostname
        self.username = username
        self.password = password
        self.port = port
        self.timeout = timeout
        self.client = None


    def connect(self):
        try:
            self.client = paramiko.SSHClient()
            self.client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
            self.client.connect(
                hostname=self.hostname,
                username=self.username,
                password=self.password,
                port=self.port,
                timeout=self.timeout
            )
            logger.info(f"Successfully connected to {self.hostname}")
        except AuthenticationException:
            logger.error("Authentication failed, please verify your credentials")
        except SSHException as sshException:
            logger.error(f"Unable to establish SSH connection: {sshException}")
        except BadHostKeyException as badHostKeyException:
            logger.error(f"Unable to verify server's host key: {badHostKeyException}")
        except Exception as e:
            logger.error(f"Exception in connecting to the server: {e}")

    def execute_command(self, command):
        if self.client is None:
            logger.error("SSH client is not connected")
            return None

        try:
            stdin, stdout, stderr = self.client.exec_command(command)
            output = stdout.read().decode('utf-8')
            error = stderr.read().decode('utf-8')
            if error:
                logger.error(f"Error executing command: {error}")
                return error
            logger.info(f"Command output: {output}")
            return output
        except Exception as e:
            logger.error(f"Exception in executing command: {e}")
            return None

    def _execute_command_thread(self, command):
        try:
            logger.info(f"Executing command: {command}")
            stdin, stdout, stderr = self.client.exec_command(command)
            output = stdout.read().decode('utf-8')
            error = stderr.read().decode('utf-8')
            if error:
                logger.error(f"Error executing command '{command}': {error}")
                return (command, error)
            logger.info(f"Command output for '{command}': {output}")
            return (command, output)
        except Exception as e:
            logger.error(f"Exception in executing command '{command}': {e}")
            return (command, str(e))

    def execute_commands(self, commands, max_workers=5):
        results = {}
        if self.client is None:
            logger.error("SSH client is not connected")
            return results

        with ThreadPoolExecutor(max_workers=max_workers) as executor:
            future_to_command = {executor.submit(self._execute_command_thread, command): command for command in commands}
            for future in as_completed(future_to_command):
                command = future_to_command[future]
                try:
                    cmd, result = future.result()
                    results[cmd] = result
                except Exception as e:
                    logger.error(f"Exception in future result for command '{command}': {e}")
                    results[command] = str(e)

        return results
    
    def save_results_to_file(self, results, file_path):
        try:
            with open(file_path, 'w', encoding='utf-8') as file:
                for command, result in results.items():
                    file.write(f"Command: {command}\nResult: {result}\n\n")
            logger.info(f"Results saved to {file_path}")
        except Exception as e:
            logger.error(f"Exception in saving results to file: {e}")    

    def close(self):
        if self.client:
            self.client.close()
            logger.info(f"Connection to {self.hostname} closed")

# 使用示例
if __name__ == "__main__":
    ssh_client = SSHClient(hostname="192.168.1.1", username="root", password="xxxx",port=22)
    ssh_client.connect()
    
    # 单个命令执行示例
    ssh_client.execute_command("ls")
    
    # 批量命令执行示例
    commands = ["ls", "pwd", "whoami"]
    results = ssh_client.execute_commands(commands)
    for cmd, result in results.items():
        print(f"Command: {cmd}\nResult: {result}\n")
    
    ssh_client.save_results_to_file(results,"ssh_command_results.txt")
    ssh_client.close()
